import backtrader as bt


class CTurtleStrategy(bt.Strategy):
    """
    海龟交易策略是众多趋势型策略中非常经典的一款，它涵盖了从入场交易、仓位管理、离场的整个过程。
    海龟交易策略采用了唐奇安通道突破作为入场和离场的信号，同时基于ATR指标动态地管理仓位，包括止盈止损以及加减仓。
    
    “唐奇安通道指标”的核心思想：通过计算一定时间范围内的最高价和最低价，形成价格的上下边界通道，用于识别价格图表上的价格趋势和波动性。
    """
    params=(("long_period", 20),
           ("short_period", 10),
           ("printlog", True))
    
    def __init__(self):
        self.order = None

        self.buy_count = 0    ## 记录买入次数
        self.buy_commision = 0.0
        self.buy_size = 0
        self.buy_price = 0

        self.close = self.datas[0].close
        self.high = self.datas[0].high
        self.low = self.datas[0].low
        
        ## 计算唐奇安通道上轨：过去20日的最高价
        self.high_line = bt.ind.Highest(self.high(-1), period=self.params.long_period, subplot=False)

        ## 计算唐奇安通道下轨：过去10日的最高价
        self.low_line = bt.ind.Lowest(self.low(-1), period=self.params.short_period, subplot=False)

        ## 计算 ATR
        self.TR = bt.ind.Max((self.high(0) - self.low(0)),
                             abs((self.high(0) - self.close(-1))),
                             abs((self.low(0) - self.close(-1))))
        self.ATR = bt.ind.MovingAverageSimple(self.TR, period=14)

        """
        self.ATR = bt.talib.ATR(self.high, self.low, self.close, timeperiod=self.params)
        """

        # 收盘价突破 ATR通道上轨，为买入信号
        self.buy_signal = bt.ind.CrossOver(self.close(0), self.high_line, subplot=False)
        
        # 收盘价跌破 ATR 通道下轨，为卖出信号
        self.sell_signal = bt.ind.CrossOver(self.close(0), self.low_line, subplot=False)
    
    def next(self):
        ## 如果还有订单在执行中，就不做新的仓位调整
        if self.order:
            return
        
        if self.buy_signal > 0 and self.buy_count == 0:
            # 进场：价格突破上轨线且空仓时
            """
            海龟交易策略使用ATR来确定每个头寸单位的规模，以便在止损时每个头寸单位损失的最大金额都是一样的，从而限制了风险。
            头寸单位的计算公式为：每单位头寸手数 = 账户余额 * 1% / ATR(N)。
            这意味着，基于N日ATR的值，交易者可以确定他们应该为每个头寸单位投入的资金量。
            """
            # self.buy_size = self.broker.getvalue() * 0.01 / (self.ATR * 1000) ## 由于ETF的波动非常小，此处人为把它放大1000倍。
            self.buy_size = self.broker.getvalue() * 0.01 / self.ATR
            self.buy_size = int(self.buy_size/100) * 100  # 1手为100的整数
            print(f"进场买入{self.buy_size}")
            self.sizer.p.stake = self.buy_size
            self.buy_count = 1
            self.order = self.buy()
            self.buy_price = self.position.price # 获取买入价格

        elif self.datas[0].close > (self.buy_price + 0.5 * self.ATR[0]) and self.buy_count > 0 and self.buy_count < 4:
            # 加仓：价格上涨了买入价的0.5的ATR且加仓次数少于3次（含）
            self.buy_size = self.broker.getvalue() * 0.01/ (self.ATR * 1000)
            self.buy_size = int(self.buy_size / 100) * 100 # 1手为100的整数
            self.sizer.p.stake = self.buy_size
            self.order = self.buy()
            self.buy_price = self.position.price # 获取买入价格
            self.buy_count += 1
            print(f"加仓买入{self.buy_size}")

        elif self.sell_signal > 0 and self.buy_count > 0:
            # 离场：价格跌破下轨线且有持仓时
            print("离场卖出")
            self.order = self.sell()
            self.buy_count = 0
        elif self.datas[0].close < (self.buy_price - 2 * self.ATR[0]) and self.buy_count > 0:
            # 止损：价格跌破买入价的2个ATR且持仓时 
            print("止损卖出")
            self.order = self.sell()
            self.buy_count = 0


    def log(self, txt, dt=None, doprint=False):
        """
        交易记录日志（默认不打印结果）
        """

        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')
    

    # 记录交易执行情况（默认不输出结果）
    def notify_order(self, order):
        pass

    # 记录交易收益情况（可省略，默认不输出结果）
    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log(f'策略收益：\n毛收益 {trade.pnl:.2f}, 净收益 {trade.pnlcomm:.2f}')

    def stop(self):
        self.log(f'(组合线：{self.p.long_period},{self.p.short_period})；\
                    期末总资金: {self.broker.getvalue():.2f}', doprint=True)

class CTurtleSizer(bt.Sizer):
    """
    由于海龟交易策略需要对仓位进行动态调整，根据ATR每次交易一单位股票，因此交易头寸需要重新设定。由CTurtleSizer类实现。
    """
    params = (('stake', 1.0),)  # stake 通常是一个浮点数，代表想要投入的资金比例

    def _getsizing(self, comminfo, cash, data, isbuy):
        if isbuy:
            # 如果是买入，返回 stake
            return self.p.stake
        else:
            # 尝试获取当前持仓
            position = self.broker.getposition(data)
            # 如果持仓为空（即没有持仓），返回0
            if not position or position.size == 0:
                return 0
            # 否则，返回当前持仓的大小（对于卖出）
            return position.size
        


#编写海龟策略
class TurtleTradingStrategy(bt.Strategy):
    params = dict(
        N1= 40, # 唐奇安通道上轨的t
        N2=30, # 唐奇安通道下轨的t
        printlog=False, 
        )
    
    def log(self, txt, dt=None,doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')
#             dt = dt or self.datas[0].datetime.date(0)
#             print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self): 
        self.order = None                   
        self.buy_count = 0 # 记录买入次数
        self.last_price = 0 # 记录买入价格
        # 准备第一个标的沪深300主力合约的close、high、low 行情数据
        self.close = self.datas[0].close
        self.high = self.datas[0].high
        self.low = self.datas[0].low
        # 计算唐奇安通道上轨：过去20日的最高价
        self.DonchianH = bt.ind.Highest(self.high(-1), period=self.p.N1, subplot=True)
        # 计算唐奇安通道下轨：过去10日的最低价
        self.DonchianL = bt.ind.Lowest(self.low(-1), period=self.p.N2, subplot=True)
        # 生成唐奇安通道上轨突破：close>DonchianH，取值为1.0；反之为 -1.0
        self.CrossoverH = bt.ind.CrossOver(self.close(0), self.DonchianH, subplot=False)
        # 生成唐奇安通道下轨突破:
        self.CrossoverL = bt.ind.CrossOver(self.close(0), self.DonchianL, subplot=False)
        # 计算 ATR
#         self.TR = bt.ind.Max((self.high(0)-self.low(0)), # 当日最高价-当日最低价
#                                     abs(self.high(0)-self.close(-1)), # abs(当日最高价−前一日收盘价)
#                                     abs(self.low(0)-self.close(-1))) # abs(当日最低价-前一日收盘价)
#         self.ATR = bt.ind.SimpleMovingAverage(self.TR, period=self.p.N1, subplot=False)
        # 计算 ATR，直接调用 talib ，使用前需要安装 python3 -m pip install TA-Lib
        self.ATR = bt.talib.ATR(self.high, self.low, self.close, timeperiod=self.p.N1, subplot=True)
#         self.ADX = bt.talib.ADX()
    
    def next(self): 
        # 如果还有订单在执行中，就不做新的仓位调整
        if self.order:
            return  
                
        # 如果当前持有多单
        if self.position.size > 0 :
            # 多单加仓:价格上涨了买入价的0.5的ATR且加仓次数少于等于3次
            if self.datas[0].close >self.last_price + 0.5*self.ATR[0] and self.buy_count <= 4:
#                 print('if self.datas[0].close >self.last_price + 0.5*self.ATR[0] and self.buy_count <= 4:')
#                 print('self.buy_count',self.buy_count)
                # 计算建仓单位：self.ATR*期货合约乘数300*保证金比例0.1
                self.buy_unit = max((self.broker.getvalue()*0.005)/(self.ATR*300*0.1),1)
                self.buy_unit = int(self.buy_unit) # 交易单位为手
                # self.sizer.p.stake = self.buy_unit
                self.order = self.buy(size=self.buy_unit)
                self.last_price = self.position.price # 获取买入价格
                self.buy_count = self.buy_count + 1
            #多单止损：当价格回落2倍ATR时止损平仓
            elif self.datas[0].close < (self.last_price - 2*self.ATR[0]):
#                 print('elif self.datas[0].close < (self.last_price - 2*self.ATR[0]):')
                self.order = self.sell(size=abs(self.position.size))
                self.buy_count = 0
            #多单止盈：当价格突破10日最低点时止盈离场 平仓
            elif self.CrossoverL < 0:
#                 print('self.CrossoverL < 0')
                self.order = self.sell(size=abs(self.position.size))
                self.buy_count = 0 
                
        # 如果当前持有空单
        elif self.position.size < 0 :
            # 空单加仓:价格小于买入价的0.5的ATR且加仓次数少于等于3次
            if self.datas[0].close<self.last_price-0.5*self.ATR[0] and self.buy_count <= 4:
#                 print('self.datas[0].close<self.last_price-0.5*self.ATR[0] and self.buy_count <= 4')
                # 计算建仓单位：self.ATR*期货合约乘数300*保证金比例0.1
                self.buy_unit = max((self.broker.getvalue()*0.005)/(self.ATR*300*0.1),1)
                self.buy_unit = int(self.buy_unit) # 交易单位为手
                # self.sizer.p.stake = self.buy_unit
                self.order = self.sell(size=self.buy_unit)
                self.last_price = self.position.price # 获取买入价格
                self.buy_count = self.buy_count + 1              
            #空单止损：当价格上涨至2倍ATR时止损平仓
            elif self.datas[0].close < (self.last_price+2*self.ATR[0]):
#                 print('self.datas[0].close < (self.last_price+2*self.ATR[0])')
                self.order = self.buy(size=abs(self.position.size))
                self.buy_count = 0
            #多单止盈：当价格突破20日最高点时止盈平仓
            elif self.CrossoverH>0:
#                 print('self.CrossoverH>0')
                self.order = self.buy(size=abs(self.position.size))
                self.buy_count = 0
                
        else: # 如果没有持仓，等待入场时机
            #入场: 价格突破上轨线且空仓时，做多
            if self.CrossoverH > 0 and self.buy_count == 0:
#                 print('if self.CrossoverH > 0 and self.buy_count == 0:')
                # 计算建仓单位：self.ATR*期货合约乘数300*保证金比例0.1
                self.buy_unit = max((self.broker.getvalue()*0.005)/(self.ATR*300*0.1),1)
                self.buy_unit = int(self.buy_unit) # 交易单位为手
                self.order = self.buy(size=self.buy_unit)
                self.last_price = self.position.price # 记录买入价格
                self.buy_count = 1  # 记录本次交易价格
            #入场: 价格跌破下轨线且空仓时，做空
            elif self.CrossoverL < 0 and self.buy_count == 0:
#                 print('self.CrossoverL < 0 and self.buy_count == 0')
                # 计算建仓单位：self.ATR*期货合约乘数300*保证金比例0.1
                self.buy_unit = max((self.broker.getvalue()*0.005)/(self.ATR*300*0.1),1)
                self.buy_unit = int(self.buy_unit) # 交易单位为手
                self.order = self.sell(size=self.buy_unit)
                self.last_price = self.position.price # 记录买入价格
                self.buy_count = 1  # 记录本次交易价格
        
    # 打印订单日志
    def notify_order(self, order):
        order_status = ['Created','Submitted','Accepted','Partial',
                        'Completed','Canceled','Expired','Margin','Rejected']
        # 未被处理的订单
        if order.status in [order.Submitted, order.Accepted]:
            self.log('ref:%.0f, name: %s, Order: %s'% (order.ref,
                                                   order.data._name,
                                                   order_status[order.status]))
            return
        # 已经处理的订单
        if order.status in [order.Partial, order.Completed]:
            if order.isbuy():
                self.log(
                        'BUY EXECUTED, status: %s, ref:%.0f, name: %s, Size: %.2f, Price: %.2f, Cost: %.2f, Comm %.2f' %
                        (order_status[order.status], # 订单状态
                         order.ref, # 订单编号
                         order.data._name, # 股票名称
                         order.executed.size, # 成交量
                         order.executed.price, # 成交价
                         order.executed.value, # 成交额
                         order.executed.comm)) # 佣金
            else: # Sell
                self.log('SELL EXECUTED, status: %s, ref:%.0f, name: %s, Size: %.2f, Price: %.2f, Cost: %.2f, Comm %.2f' %
                            (order_status[order.status],
                             order.ref,
                             order.data._name,
                             order.executed.size,
                             order.executed.price,
                             order.executed.value,
                             order.executed.comm))
                    
        elif order.status in [order.Canceled, order.Margin, order.Rejected, order.Expired]:
            # 订单未完成
            self.log('ref:%.0f, name: %s, status: %s'% (
                order.ref, order.data._name, order_status[order.status]))
            
        self.order = None
        
    def notify_trade(self, trade):
        # 交易刚打开时
        if trade.justopened:
            self.log('Trade Opened, name: %s, Size: %.2f,Price: %.2f' % (
                    trade.getdataname(), trade.size, trade.price))
        # 交易结束
        elif trade.isclosed:
            self.log('Trade Closed, name: %s, GROSS %.2f, NET %.2f, Comm %.2f' %(
            trade.getdataname(), trade.pnl, trade.pnlcomm, trade.commission))
        # 更新交易状态
        else:
            self.log('Trade Updated, name: %s, Size: %.2f,Price: %.2f' % (
                    trade.getdataname(), trade.size, trade.price))
    def stop(self):
      
        self.log(f'(组合线：{self.p.N1},{self.p.N2})； \
        期末总资金: {self.broker.getvalue():.2f}', doprint=True)

